import torch
from torch import nn
import matplotlib.pyplot as plt

# -------------------准备模型训练数据--------------------
# 加载保存的数据
features = torch.load("features.pth", weights_only=False)
# 输入
train_features = features[:, :2]
# 输出（真实的类别）
train_labels = features[:, 2:]

# -------------------分类原始数据可视化--------------------

fig, axes = plt.subplots(1, 2)


def draw_points(features, labels, num):
    for i, item in enumerate(features):
        x, y = item
        label = labels[i][0]
        color = "ro"
        if label == 1:
            color = "bo"
        axes[num].plot(x, y, color)


draw_points(train_features, train_labels, 0)
# -------------------准备模型--------------------
model = nn.Sequential(
    nn.Linear(2, 10),
    nn.Tanh(),
    nn.Linear(10, 1),
    nn.Sigmoid()  # sigmoid 函数的取值范围 0~1
)
# -------------------损失函数--------------------
criterion = nn.MSELoss()
# -------------------优化器（梯度下降）--------------------
sgd = torch.optim.SGD(model.parameters(), lr=0.05)
# -------------------训练数据--------------------
epochs = 500
for epoch in range(epochs):
    sgd.zero_grad()
    predict_labels = model(train_features)
    loss = criterion(predict_labels, train_labels)
    loss.backward()
    sgd.step()

    print(f"epoch {epoch + 1}/{epochs} -- loss:{loss.item():.4f}")

# ------------------测试整个模型的准确性--------------------
model.eval()  # 开启测试模式
# 设计一个数据集，作为此次模型测试的数据
x = torch.linspace(0, 2, 20)
y = torch.linspace(0, 2, 20)
x, y = torch.meshgrid([x, y], indexing="ij")
x = x.reshape(400, 1)
y = y.reshape(400, 1)
# 测试使用的xy集合
test_features = torch.concatenate([x, y], dim=-1)
# 直接使用训练好的模型预测它们的分类
test_labels = model(test_features)
# sigmoid 的取值范围是0~1之间所以，我们可以认为接近0的区域都是0，接近1的区域都是1
test_labels[test_labels >= 0.5] = 1
test_labels[test_labels < 0.5] = 0
# -----------------绘制测试模型---------------------
draw_points(test_features, test_labels, 1)

plt.show()
